package com.deepseek.project.tool;

import com.deepseek.project.constant.Constants;
import com.deepseek.project.service.ITalkInfoService;
import com.google.gson.reflect.TypeToken;
import okhttp3.*;
import okio.BufferedSource;
import org.jetbrains.annotations.NotNull;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.IOException;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;


/**
 * author:hulei
 * 调用本地部署的ollama+DeepSeek模型
 */
@Component
public class DeepSeekLocalTool extends AbstractChatTool {

    private static final String URL = "http://127.0.0.1:11434/api/chat";

    private static final String MODEL = "deepseek-r1:32b";

    public DeepSeekLocalTool(ITalkInfoService italkInfoService, List<Map<String, String>> conversationHistory) {
        super(italkInfoService, conversationHistory);
    }

    @Override
    protected Request buildRequest() {
        Map<String, Object> data = new HashMap<>();
        data.put("model", MODEL);
        data.put("messages", Collections.unmodifiableList(getConversationHistory()));
        data.put("stream", true);

        MediaType JSON = MediaType.get("application/json; charset=utf-8");
        RequestBody body = RequestBody.create(gson.toJson(data), JSON);
        return new Request.Builder()
                .url(URL)
                .post(body).build();
    }

    /**
     * 请求DeepSeek深度搜索消息
     *
     * @param session WebSocket会话WebSocketSession
     * @param talkInfoId 会话框id
     */
    @Override
    public void requestDeepSeekMessage(WebSocketSession session, String talkInfoId) {
        Request request = buildRequest();
        client.newCall(request).enqueue(new Callback() {
            @Override
            public void onFailure(@NotNull Call call, @NotNull IOException e) {
                handleError(session, e);
            }

            @Override
            public void onResponse(@NotNull Call call, @NotNull Response response) {
                try (ResponseBody body = response.body()) {
                    if (body == null) return;
                    processStreamResponse(body.source(), session, talkInfoId);
                    // 发送流式响应结束标记
                    sendEndMarker(session);
                } catch (IOException e) {
                    handleError(session, e);
                }
            }
        });
    }


    public void processStreamResponse(BufferedSource source, WebSocketSession session, String talkInfoId) throws IOException {
        String line;
        while (!stopSending && (line = source.readUtf8Line()) != null) {
            Map<String, Object> responseMap = parseResponseLine(line);
            String content = getContent(responseMap);
            Map<String, String> assistantMessageMap = Map.of("role", "assistant", "content", content);
            synchronized (conversationHistory) {
                conversationHistory.add(assistantMessageMap);
            }
            if (!talkInfoId.startsWith(Constants.TEMP_TALK_INFO_ID_PREFIX)) {
                addMessageToCache(assistantMessageMap, talkInfoId);
            }
            if (!stopSending) {
                session.sendMessage(new TextMessage(content));
            }
        }
    }

    public String getContent(Map<String, Object> responseMap) {
        @SuppressWarnings("unchecked")
        Map<String, Object> messageMap = (Map<String, Object>) responseMap.get("message");
        return (String) messageMap.get("content");
    }

    public Map<String, Object> parseResponseLine(String line) {
        return gson.fromJson(line, new TypeToken<Map<String, Object>>() {
        }.getType());
    }

}
